"""Schema 提取模块"""
from sqlalchemy import text, inspect
from typing import List, Dict


class SchemaExtractor:
    """数据库 Schema 提取器"""
    
    def __init__(self, engine, db_type):
        """
        初始化 Schema 提取器
        
        Args:
            engine: SQLAlchemy engine 对象
            db_type: 数据库类型 ('mysql' 或 'postgresql')
        """
        self.engine = engine
        self.db_type = db_type
        self.inspector = inspect(engine)
    
    def get_tables(self) -> List[str]:
        """
        获取数据库中所有表名
        
        Returns:
            表名列表
        """
        try:
            tables = self.inspector.get_table_names()
            return tables
        except Exception as e:
            raise Exception(f"获取表列表失败: {str(e)}")
    
    def get_table_columns(self, table_name: str) -> List[Dict]:
        """
        获取指定表的所有字段信息
        
        Args:
            table_name: 表名
            
        Returns:
            字段信息列表，每个字段包含：
            - column_name: 字段名
            - data_type: 数据类型
            - length: 长度/精度
            - is_nullable: 是否允许空
            - default_value: 默认值
            - comment: 备注
        """
        try:
            # 直接从 information_schema 查询，获取更准确的信息
            if self.db_type == 'mysql':
                return self._get_mysql_columns(table_name)
            elif self.db_type == 'postgresql':
                return self._get_postgresql_columns(table_name)
            else:
                raise ValueError(f"不支持的数据库类型: {self.db_type}")
            
        except Exception as e:
            raise Exception(f"获取表 '{table_name}' 的字段信息失败: {str(e)}")
    
    def _get_mysql_columns(self, table_name: str) -> List[Dict]:
        """获取 MySQL 表的字段信息"""
        query = text("""
            SELECT 
                COLUMN_NAME,
                DATA_TYPE,
                CHARACTER_MAXIMUM_LENGTH,
                NUMERIC_PRECISION,
                NUMERIC_SCALE,
                IS_NULLABLE,
                COLUMN_DEFAULT,
                COLUMN_COMMENT
            FROM INFORMATION_SCHEMA.COLUMNS
            WHERE TABLE_SCHEMA = DATABASE()
            AND TABLE_NAME = :table_name
            ORDER BY ORDINAL_POSITION
        """)
        
        result = []
        with self.engine.connect() as conn:
            rows = conn.execute(query, {'table_name': table_name}).fetchall()
            
            for row in rows:
                data_type = row[1].upper()
                length = None
                
                # 处理长度/精度
                if row[2] is not None:  # CHARACTER_MAXIMUM_LENGTH
                    length = str(row[2])
                elif row[3] is not None:  # NUMERIC_PRECISION
                    if row[4] is not None and row[4] > 0:  # NUMERIC_SCALE
                        length = f"{row[3]},{row[4]}"
                    else:
                        length = str(row[3])
                
                column_info = {
                    'column_name': row[0],
                    'data_type': data_type,
                    'length': length,
                    'is_nullable': row[5] == 'YES',
                    'default_value': str(row[6]) if row[6] is not None else '',
                    'comment': row[7] or ''
                }
                result.append(column_info)
        
        return result
    
    def _get_postgresql_columns(self, table_name: str) -> List[Dict]:
        """获取 PostgreSQL 表的字段信息"""
        query = text("""
            SELECT 
                c.column_name,
                c.data_type,
                c.character_maximum_length,
                c.numeric_precision,
                c.numeric_scale,
                c.is_nullable,
                c.column_default,
                pgd.description as column_comment
            FROM information_schema.columns c
            LEFT JOIN pg_catalog.pg_statio_all_tables as st on c.table_schema=st.schemaname and c.table_name=st.relname
            LEFT JOIN pg_catalog.pg_description pgd on (pgd.objoid=st.relid and pgd.objsubid=c.ordinal_position)
            WHERE c.table_schema = 'public'
            AND c.table_name = :table_name
            ORDER BY c.ordinal_position
        """)
        
        result = []
        with self.engine.connect() as conn:
            rows = conn.execute(query, {'table_name': table_name}).fetchall()
            
            for row in rows:
                data_type = row[1].upper()
                length = None
                
                # 处理长度/精度
                if row[2] is not None:  # character_maximum_length
                    length = str(row[2])
                elif row[3] is not None:  # numeric_precision
                    if row[4] is not None and row[4] > 0:  # numeric_scale
                        length = f"{row[3]},{row[4]}"
                    else:
                        length = str(row[3])
                
                # 获取注释
                comment = row[7] if row[7] else ''
                
                column_info = {
                    'column_name': row[0],
                    'data_type': data_type,
                    'length': length,
                    'is_nullable': row[5] == 'YES',
                    'default_value': str(row[6]) if row[6] is not None else '',
                    'comment': comment
                }
                result.append(column_info)
        
        return result
    
    def get_table_schema(self) -> Dict[str, List[Dict]]:
        """
        获取整个数据库的 Schema 信息
        
        Returns:
            字典，key 为表名，value 为字段信息列表
        """
        schema = {}
        tables = self.get_tables()
        
        for table_name in tables:
            schema[table_name] = self.get_table_columns(table_name)
        
        return schema

